import glob
import json
import os
from os.path import join
import pickle
import re
from ast import parse
import sys
sys.path.append("../")
import numpy as np
import trimesh
from scipy.spatial.distance import cdist
from termcolor import cprint
from tqdm import tqdm

from tink.cal_contact_info import load_pointcloud, to_pointcloud

import DeepSDF_OakInk.deep_sdf.workspace as ws
import torch
from DeepSDF_OakInk.deep_sdf.mesh import create_mesh
import time
import open3d as o3d
from tqdm import tqdm


def gen_interpolate(data, source, target, interpolate=10):
    handle_list = []

    handle_list.append((source, target))

    print(len(handle_list))

    for idx, (soruce_idx, target_idx) in enumerate(handle_list):

        cprint(f"handling {soruce_idx} -> {target_idx}", "blue")

        specs = json.load(open(os.path.join(data, "specs.json"), "r"))
        arch = __import__("DeepSDF_OakInk.networks." + specs["NetworkArch"], fromlist=["Decoder"])
        latent_size = specs["CodeLength"]
        decoder = arch.Decoder(latent_size, **specs["NetworkSpecs"])
        decoder = torch.nn.DataParallel(decoder)
        saved_model_state = torch.load(os.path.join(data, "network", ws.model_params_subdir, "latest" + ".pth"))
        saved_model_epoch = saved_model_state["epoch"]
        decoder.load_state_dict(saved_model_state["model_state_dict"])

        decoder = decoder.module.cuda()

        sdf_code_path = os.path.join(data, ws.reconstructions_subdir, ws.reconstruction_codes_subdir)
        sdf_mesh_path = os.path.join(data, ws.reconstructions_subdir, ws.reconstruction_meshes_subdir)
        source = f"{soruce_idx}.pth"
        target = f"{target_idx}.pth"
        interpolate_path = os.path.join(data, "interpolate", f"{soruce_idx}-{target_idx}")
        if os.path.exists(interpolate_path) and len(os.listdir(interpolate_path)) == interpolate:
            continue

        os.makedirs(interpolate_path, exist_ok=True)

        k = interpolate

        with torch.no_grad():
            latent_source = torch.load(os.path.join(sdf_code_path, source))
            latent_target = torch.load(os.path.join(sdf_code_path, target))
            if not os.path.isfile(os.path.join(sdf_mesh_path, source.replace(".pth", ".ply"))):
                create_mesh(
                    decoder,
                    latent_source[0].type(torch.FloatTensor).cuda(),
                    os.path.join(sdf_mesh_path, source),
                    N=256,
                    max_batch=int(2 ** 18),
                )
            if not os.path.isfile(os.path.join(sdf_mesh_path, target.replace(".pth", ".ply"))):
                create_mesh(
                    decoder,
                    latent_target[0].type(torch.FloatTensor).cuda(),
                    os.path.join(sdf_mesh_path, target),
                    N=256,
                    max_batch=int(2 ** 18),
                )

            latent_source = latent_source.cpu().numpy()[0][0]
            latent_target = latent_target.cpu().numpy()[0][0]

        inter_list = []
        for i in range(256):
            y = np.array([latent_source[i], latent_target[i]])
            x = np.array([0, k + 1])
            xl = np.arange(1, k + 1)
            interp = np.interp(xl, x, y)
            inter_list.append(interp)
        interpolation = np.vstack(inter_list)

        with torch.no_grad():
            for i in range(k):
                interp_obj = os.path.join(interpolate_path, f"interp{i + 1:02d}")

                latent = torch.from_numpy(interpolation[:, i].T)
                create_mesh(
                    decoder,
                    latent.type(torch.FloatTensor).cuda(),
                    interp_obj,
                    N=256,
                    max_batch=int(2 ** 18),
                )


def get_obj_path(oid, use_downsample=True, key="align"):
    obj_suffix_path = "align_ds" if use_downsample else "align"
    real_meta = json.load(open("/home/liuyun/HHO-dataset/data_processing/Tink/DeepSDF_OakInk/data/meta/object_id.json", "r"))
    virtual_meta = json.load(open("/home/liuyun/HHO-dataset/data_processing/Tink/DeepSDF_OakInk/data/meta/virtual_object_id.json", "r"))
    if oid in real_meta:
        obj_name = real_meta[oid]["name"]
        obj_path = "/home/liuyun/HHO-dataset/data_processing/Tink/DeepSDF_OakInk/data/OakInkObjects"
    else:
        obj_name = virtual_meta[oid]["name"]
        obj_path = "/home/liuyun/HHO-dataset/data_processing/Tink/DeepSDF_OakInk/data/OakInkVirtualObjects"
    obj_mesh_path = list(
        glob.glob(os.path.join(obj_path, obj_name, obj_suffix_path, "*.obj"))
        + glob.glob(os.path.join(obj_path, obj_name, obj_suffix_path, "*.ply"))
    )
    if len(obj_mesh_path) > 1:
        obj_mesh_path = [p for p in obj_mesh_path if key in os.path.split(p)[1]]
    assert len(obj_mesh_path) == 1
    return obj_mesh_path[0]

def cal_closest_idx(source, target):
    region = np.max([np.abs(source.min()), source.max(), np.abs(target.min()), target.max()])
    k = 6
    cell_size = region * 2 / k
    res = np.zeros(len(target), dtype=int)
    for x in range(k):
        for y in range(k):
            for z in range(k):

                def select_region(d, i):
                    return (target[:, d] > i * cell_size - region - 1e-6) & (
                        target[:, d] < (i + 1) * cell_size - region + 1e-6
                    )

                def search_region(d, i):
                    return (source[:, d] > (i - 1) * cell_size - region - 1e-6) & (
                        source[:, d] < (i + 2) * cell_size - region + 1e-6
                    )

                target_filter = select_region(0, x) & select_region(1, y) & select_region(2, z)
                source_filter = search_region(0, x) & search_region(1, y) & search_region(2, z)
                target_cell = target[target_filter]
                source_cell = source[source_filter]
                if len(target_cell) == 0 or len(source_cell) == 0:
                    continue
                res_cell = np.argwhere(source_filter)[np.argmin(cdist(target_cell, source_cell), axis=1)].ravel()
                res[target_filter] = res_cell
    return res

def transform_contact(data, source, target, origin_mesh, contact=None, thresh=0.003, source_contact=None):
    '''
    contact: list of contact vertice id in source mesh, N 
    '''

    interpolate_path = os.path.join(data, "interpolate", f"{source}-{target}")

    inter_paths = os.listdir(interpolate_path)
    inter_paths.sort()

    inter_pc = [load_pointcloud(os.path.join(interpolate_path, p)) for p in inter_paths]

    rescale = pickle.load(open(os.path.join(data, "rescale.pkl"), "rb"))
    rescale = rescale["max_norm"] * rescale["scale"]

    source_mesh = trimesh.load(get_obj_path(source), process=False, force="mesh", skip_materials=True)

    target_mesh = trimesh.load(get_obj_path(target), process=False, force="mesh", skip_materials=True)
    # source_mesh = trimesh.load_mesh(get_obj_path(source), )
    # target_mesh = trimesh.load_mesh(get_obj_path(target))

    source_mesh.vertices = (
        source_mesh.vertices - (np.array(source_mesh.vertices).max(0) + np.array(source_mesh.vertices).min(0)) / 2
    )  # center
    source_rescale_x = np.array(source_mesh.vertices)[:, 0].max()
    source_rescale_y = np.array(source_mesh.vertices)[:, 1].max()
    source_rescale_z = np.array(source_mesh.vertices)[:, 2].max()
    source_rescale = np.array([source_rescale_x, source_rescale_y, source_rescale_z])
    # for i in range(3):
    #     source_mesh.vertices[:, i] = np.array(source_mesh.vertices[:, i]) / source_rescale[i]
    source_mesh.vertices = np.array(source_mesh.vertices) / rescale
    interpolate_path = os.path.join(data, "interpolate", f"{source}-{target}")

    inter_paths = os.listdir(interpolate_path)
    inter_paths.sort()

    inter_pc = [load_pointcloud(os.path.join(interpolate_path, p)) for p in inter_paths]

    # for p in inter_pc:
    #     points = np.asarray(p.points)
    #     point_rescale_x = points[:, 0].max()
    #     point_rescale_y = points[:, 1].max()
    #     point_rescale_z = points[:, 2].max()
    #     point_rescale = [point_rescale_x, point_rescale_y, point_rescale_z]
    #     for i in range(3):
    #         points[:, i] =  points[:, i] / point_rescale[i]
    #     p.points = o3d.utility.Vector3dVector(points)

    target_mesh.vertices = (
        target_mesh.vertices - (np.array(target_mesh.vertices).max(0) + np.array(target_mesh.vertices).min(0)) / 2
    )  # center
    target_rescale_x = np.array(target_mesh.vertices)[:, 0].max()
    target_rescale_y = np.array(target_mesh.vertices)[:, 1].max()
    target_rescale_z = np.array(target_mesh.vertices)[:, 2].max()

    target_rescale = np.array([target_rescale_x, target_rescale_y, target_rescale_z])
    # for i in range(3):
    #     target_mesh.vertices[:, i] = np.array(target_mesh.vertices[:, i]) / target_rescale[i]
    target_mesh.vertices = target_mesh.vertices / rescale

    # transform
    # hho_mesh = origin_mesh
    # hho_mesh.vertices = (
    #     hho_mesh.vertices - (np.array(hho_mesh.vertices).max(0) + np.array(hho_mesh.vertices).min(0)) / 2
    # )  # center
    # hho_mesh.vertices = hho_mesh.vertices / rescale

    source_pointcloud = to_pointcloud(source_mesh)
    target_pointcloud = to_pointcloud(target_mesh)

    # hho_points = [] # N x 3
    # assert contact != None
    # for point in contact:
    #     hho_points.append(hho_mesh.vertices[point])

    source_vert = np.asarray(source_mesh.vertices)
    # source_vert_len = len(source_mesh.vertices)

    source_contact_area = contact # N x index_num
    # for hho_point in tqdm(hho_points):
    #     for i in range(source_vert_len):
    #         if np.linalg.norm(source_vert[i] - hho_point) < thresh:
    #             source_contact_area.append(i)
    obj_pointclouds = [source_vert] + [ip.points for ip in inter_pc] + [target_pointcloud.points]
    target_contact_area = source_contact_area
    for i, (source, target) in enumerate(zip(obj_pointclouds[:-1], obj_pointclouds[1:])):
        source_p = np.array(source)
        target_p = np.array(target)
        source_rescale = np.array([source_p[:, 0].max(), source_p[:, 1].max(), source_p[:, 2].max()])
        target_rescale = np.array([target_p[:, 0].max(), target_p[:, 1].max(), target_p[:, 2].max()])
        for j in range(3):
            source_p[:, j] = source_p[:, j] / source_rescale[j]
            target_p[:, j] = target_p[:, j] / target_rescale[j]
        target_idx = cal_closest_idx(target_p, source_p)
        target_contact_area = target_idx[target_contact_area]

    return source_contact_area, target_contact_area

def vis(source_contact, source, target_contact, target, person, vis_save_file):
    # 可视化 contact 保存
    # source = o3d.io.read_point_cloud(get_obj_path(source))
    source = trimesh.load(get_obj_path(source), process=False, force="mesh", skip_materials=True)
    source_pcd = o3d.geometry.PointCloud()
    source_pcd.points = o3d.utility.Vector3dVector(np.asarray(source.vertices))
    source_pcd.paint_uniform_color([0, 0, 1])
    color = np.asarray(source_pcd.colors)
    color[source_contact] = [1, 0, 0]
    source_pcd.colors = o3d.utility.Vector3dVector(color)
    o3d.io.write_point_cloud(join(vis_save_file, f"{person}_source_contact.ply"), source_pcd)

    print(get_obj_path(target))
    # target = o3d.io.read_triangle_mesh(get_obj_path(target))
    target_mesh = trimesh.load(get_obj_path(target), process=False, force="mesh", skip_materials=True)
    
    target_pc = to_pointcloud(target_mesh)
    target_pc.paint_uniform_color([0, 0, 1])
    color = np.asarray(target_pc.colors)
    color[target_contact] = [1, 0, 0]
    target_pc.colors = o3d.utility.Vector3dVector(color)
    o3d.io.write_point_cloud(join(vis_save_file, f"{person}_target_contact.ply"), target_pc)


def vis_pcd(contact, pcd, write_path):
    pcd.paint_uniform_color([0, 0, 1])
    color = np.asarray(pcd.colors)
    color[contact["person1"]] = [1, 0, 0]
    color[contact["person2"]] = [0, 1, 0]
    pcd.colors = o3d.utility.Vector3dVector(color)
    o3d.io.write_point_cloud(write_path, pcd)


def tranfer_contact_to_new_obj(data, source, target, origin_mesh, type, source_contact, person=None, visualize=False, interpolate=10, vis_save_file=None):
    source_obj = source
    target_obj = target
    data = os.path.join(data, type)
    gen_interpolate(data, source_obj, target_obj, interpolate=interpolate)
    source_contact, target_contact = transform_contact(data, source_obj, target_obj, origin_mesh, source_contact, thresh=0.01, source_contact=source_contact)
    if visualize:
        print("visualizing...")
        vis(source_contact, source_obj, target_contact, target_obj, person, vis_save_file=vis_save_file)

    return source_contact, target_contact

    



if __name__ == '__main__':
    data = "/home/liuyun/HHO-dataset/data_processing/Tink/DeepSDF_OakInk/data/sdf/"
    source = "h01chair9"
    target = "s25112"
    hho_file = "/share/datasets/HHO_object_dataset_final/椅子/椅子9/chair9_m.obj"

    hho_mesh = trimesh.load_mesh(hho_file)
    vert = np.asarray(hho_mesh.vertices)
    source_contact = [np.argmax(vert[:, 1])] + [np.argmax(vert[:, 0])]
    print(source_contact)
    source_contact, target_contact = tranfer_contact_to_new_obj(data, source, target, hho_file, "chair", source_contact, visualize=True)
    # gen_interpolate(data, source, target)
    # source_contact = range(20)
    # source_contact, target_contact = transform_contact(data, source, target, "/share/datasets/HHO_object_dataset_final/椅子/椅子5/chair5_m.obj", source_contact, thresh=0.01)
    # vis(source_contact, source, target_contact, target)




    # target_1 = trimesh.load(get_obj_path(target), process=False, force="mesh", skip_materials=True)
    # target_2 = trimesh.load(get_obj_path(target), force="mesh", skip_materials=True)

    # target_3 = to_pointcloud(target_1)
    
    # print(target_1.vertices.shape)
    # print(target_2.vertices.shape)
    # target_1_pc = o3d.geometry.PointCloud()
    # target_1_pc.points = o3d.utility.Vector3dVector(np.asarray(target_1.vertices))

    # target_2_pc = o3d.geometry.PointCloud()
    # target_2_pc.points = o3d.utility.Vector3dVector(np.asarray(target_2.vertices))


    # o3d.io.write_point_cloud("target_1.ply", target_1_pc)
    # o3d.io.write_point_cloud("target_3.ply", target_3)